from PyQt5.QtWidgets import QWidget, QMenu, QAction, QInputDialog
from PyQt5.QtGui import QPainter, QPen, QFont, QCursor, QPixmap
from PyQt5.QtCore import Qt, QRectF, QPointF
from typing import List, Dict, Optional
from sql_parser import Table, Column
import math

class ERShape:
    def __init__(self, x: float, y: float, width: float, height: float):
        self.rect = QRectF(x - width/2, y - height/2, width, height)
        self.is_selected = False
        self.is_dragging = False
        self.drag_offset = QPointF()

    def contains(self, point: QPointF) -> bool:
        return self.rect.contains(point)

    def start_drag(self, point: QPointF):
        self.is_dragging = True
        self.drag_offset = QPointF(point.x() - self.rect.center().x(),
                                 point.y() - self.rect.center().y())

    def drag_to(self, point: QPointF):
        if self.is_dragging:
            new_center = QPointF(point.x() - self.drag_offset.x(),
                               point.y() - self.drag_offset.y())
            self.rect.moveCenter(new_center)

    def stop_drag(self):
        self.is_dragging = False

class Entity(ERShape):
    def __init__(self, name: str, x: float, y: float):
        super().__init__(x, y, 120, 60)
        self.name = name

class Attribute(ERShape):
    def __init__(self, name: str, x: float, y: float, is_primary: bool = False, is_foreign: bool = False):
        super().__init__(x, y, 100, 40)
        self.name = name
        self.is_primary = is_primary
        self.is_foreign = is_foreign

class ERGenerator:
    def __init__(self, widget: QWidget):
        self.widget = widget
        self.tables: Dict[str, Table] = {}
        self.current_table = None
        self.entity: Optional[Entity] = None
        self.attributes: List[Attribute] = []
        self.selected_shape = None
        self.widget.setGenerator(self)
        
        # 设置鼠标跟踪
        self.widget.setMouseTracking(True)
        
        # 创建右键菜单
        self.context_menu = QMenu(self.widget)
        self.add_attr_action = QAction("添加属性", self.widget)
        self.edit_action = QAction("编辑", self.widget)
        self.delete_action = QAction("删除", self.widget)
        
        self.context_menu.addAction(self.add_attr_action)
        self.context_menu.addAction(self.edit_action)
        self.context_menu.addAction(self.delete_action)
        
        # 连信号
        self.add_attr_action.triggered.connect(self.add_attribute)
        self.edit_action.triggered.connect(self.edit_shape)
        self.delete_action.triggered.connect(self.delete_shape)
        
        self.scale = 1.0  # 缩放比例
        self.min_scale = 0.5  # 最小缩放
        self.max_scale = 2.0  # 最大缩放
        self.offset_x = 0  # 添加平移偏移量
        self.offset_y = 0
        self.is_panning = False  # 是否正在平移
        self.last_mouse_pos = None

    def generate(self, tables: List[Table]):
        self.tables = {table.name: table for table in tables}
        if tables:
            self.show_table(tables[0].name)

    def show_table(self, table_name: str):
        if table_name in self.tables:
            self.current_table = self.tables[table_name]
            # 清除旧的实体和属性
            self.entity = None
            self.attributes = []
            self._calculate_positions()
            self.widget.update()

    def _calculate_positions(self):
        if not self.current_table:
            return
            
        width = self.widget.width()
        height = self.widget.height()
        
        # 如果实体不存在才创建新的
        if not self.entity:
            center_x = width / 2
            center_y = height / 2
            self.entity = Entity(self.current_table.comment or self.current_table.name,
                               center_x, center_y)
            self.attributes = []
            
            # 创建属性
            attrs = self.current_table.columns
            num_attrs = len(attrs)
            
            if num_attrs > 0:
                radius = min(width, height) * 0.2
                
                # 计算每个属性的角度，从正上方开始，顺时针均匀分布
                for i, attr in enumerate(attrs):
                    # 计算角度：360度平均分配
                    angle = -math.pi/2 + (2 * math.pi * i) / num_attrs
                    
                    # 计算位置
                    x = center_x + radius * math.cos(angle)
                    y = center_y + radius * math.sin(angle)
                    
                    attribute = Attribute(attr.comment or attr.name, x, y,
                                        attr.is_primary, attr.is_foreign)
                    self.attributes.append(attribute)
        
        # 应用缩放和平移
        for shape in [self.entity] + self.attributes:
            if shape:
                # 保存原始中心点
                old_center = shape.rect.center()
                
                # 应用缩放
                if isinstance(shape, Entity):
                    shape.rect.setWidth(120 * self.scale)
                    shape.rect.setHeight(60 * self.scale)
                else:
                    shape.rect.setWidth(100 * self.scale)
                    shape.rect.setHeight(40 * self.scale)
                
                # 应用平移
                new_x = old_center.x() + self.offset_x
                new_y = old_center.y() + self.offset_y
                shape.rect.moveCenter(QPointF(new_x, new_y))

    def handle_mouse_press(self, event):
        if event.button() == Qt.LeftButton:
            point = event.pos()
            # 检查是否点击了任何图形
            self.selected_shape = None
            
            # 转换点坐标以匹配缩放后的位置
            scaled_point = QPointF(
                (point.x() - self.widget.width()/2) / self.scale + self.widget.width()/2,
                (point.y() - self.widget.height()/2) / self.scale + self.widget.height()/2
            )
            
            if self.entity and self.entity.contains(scaled_point):
                self.selected_shape = self.entity
            else:
                for attr in self.attributes:
                    if attr.contains(scaled_point):
                        self.selected_shape = attr
                        break
            
            if self.selected_shape:
                self.selected_shape.start_drag(scaled_point)
                self.widget.update()
        
        elif event.button() == Qt.RightButton:
            point = event.pos()
            # 显示右键菜单
            self.show_context_menu(point)

    def handle_mouse_move(self, event):
        if self.selected_shape and self.selected_shape.is_dragging:
            # 转换点坐标以匹配缩放后的位置
            scaled_point = QPointF(
                (event.pos().x() - self.widget.width()/2) / self.scale + self.widget.width()/2,
                (event.pos().y() - self.widget.height()/2) / self.scale + self.widget.height()/2
            )
            self.selected_shape.drag_to(scaled_point)
            self.widget.update()

    def handle_mouse_release(self, event):
        if event.button() == Qt.LeftButton and self.selected_shape:
            self.selected_shape.stop_drag()

    def show_context_menu(self, point):
        self.selected_shape = None
        if self.entity and self.entity.contains(point):
            self.selected_shape = self.entity
        else:
            for attr in self.attributes:
                if attr.contains(point):
                    self.selected_shape = attr
                    break
        
        if self.selected_shape:
            self.context_menu.exec_(self.widget.mapToGlobal(point))

    def add_attribute(self):
        name, ok = QInputDialog.getText(self.widget, "添加属性", "属性名称:")
        if ok and name:
            # 在当前位置添加新属性
            pos = self.widget.mapFromGlobal(QCursor.pos())
            new_attr = Attribute(name, pos.x(), pos.y())
            self.attributes.append(new_attr)
            self.widget.update()

    def edit_shape(self):
        if self.selected_shape:
            new_name, ok = QInputDialog.getText(self.widget, "编辑",
                                              "名称:", text=self.selected_shape.name)
            if ok and new_name:
                self.selected_shape.name = new_name
                self.widget.update()

    def delete_shape(self):
        if self.selected_shape in self.attributes:
            # 删除属性
            self.attributes.remove(self.selected_shape)
            self.selected_shape = None
            
            # 重新计算所有属性的位置
            if self.attributes:
                width = self.widget.width()
                height = self.widget.height()
                center_x = width / 2
                center_y = height / 2
                radius = min(width, height) * 0.2
                
                # 重新计算每个属性的位置
                num_attrs = len(self.attributes)
                for i, attr in enumerate(self.attributes):
                    # 计算新的角度和位置
                    angle = -math.pi/2 + (2 * math.pi * i) / num_attrs
                    x = center_x + radius * math.cos(angle)
                    y = center_y + radius * math.sin(angle)
                    
                    # 更新属性位置
                    attr.rect.moveCenter(QPointF(x, y))
            
            # 更新显示
            self.widget.update()

    def _find_intersection_point(self, start_x: float, start_y: float, 
                               end_x: float, end_y: float, 
                               rect: QRectF, is_ellipse: bool = False) -> tuple:
        """计算线段与矩形/椭圆的交点"""
        center_x = rect.center().x()
        center_y = rect.center().y()
        
        # 计算方向向量
        dx = end_x - start_x
        dy = end_y - start_y
        
        if abs(dx) < 0.001 and abs(dy) < 0.001:
            return start_x, start_y
        
        if is_ellipse:
            # 椭圆的半宽和半高
            a = rect.width() / 2
            b = rect.height() / 2
            
            # 计算角度
            angle = math.atan2(dy, dx)
            
            # 使用椭圆参数方程
            x = center_x - a * math.cos(angle)
            y = center_y - b * math.sin(angle)
        else:
            # 矩形的情况
            half_width = rect.width() / 2
            half_height = rect.height() / 2
            
            # 计算从起点到终点的线段与矩形四边的交点
            intersections = []
            
            # 检查与上边的交点
            if abs(dy) > 0.001:
                y = rect.top()
                x = start_x + (dx/dy) * (y - start_y)
                if rect.left() <= x <= rect.right() and (y - start_y) / dy > 0:
                    intersections.append((x, y, abs(x - start_x) + abs(y - start_y)))
            
            # 检查与下边的交点
            if abs(dy) > 0.001:
                y = rect.bottom()
                x = start_x + (dx/dy) * (y - start_y)
                if rect.left() <= x <= rect.right() and (y - start_y) / dy > 0:
                    intersections.append((x, y, abs(x - start_x) + abs(y - start_y)))
            
            # 检查与左边的交点
            if abs(dx) > 0.001:
                x = rect.left()
                y = start_y + (dy/dx) * (x - start_x)
                if rect.top() <= y <= rect.bottom() and (x - start_x) / dx > 0:
                    intersections.append((x, y, abs(x - start_x) + abs(y - start_y)))
            
            # 检查与右边的交点
            if abs(dx) > 0.001:
                x = rect.right()
                y = start_y + (dy/dx) * (x - start_x)
                if rect.top() <= y <= rect.bottom() and (x - start_x) / dx > 0:
                    intersections.append((x, y, abs(x - start_x) + abs(y - start_y)))
            
            # 如果找到交点，选择最近的一个
            if intersections:
                # 按距离排序
                intersections.sort(key=lambda p: p[2])
                return intersections[0][0], intersections[0][1]
            
            # 如果没有找到交点，返回起点
            return start_x, start_y
        
        return x, y

    def paint_for_export(self, painter: QPainter, min_x: float, min_y: float, width: float, height: float):
        if not self.entity:
            return
            
        painter.setRenderHint(QPainter.Antialiasing)
        
        # 设置字体
        entity_font = QFont("Arial", 12, QFont.Bold)
        attr_font = QFont("Arial", 10)
        
        # 绘制实体
        painter.setPen(QPen(Qt.black, 2))
        painter.drawRect(self.entity.rect)
        painter.setFont(entity_font)
        painter.drawText(self.entity.rect, Qt.AlignCenter, self.entity.name)
        
        # 绘制连线和属性
        for attr in self.attributes:
            # 获取实体和属性的中心点
            entity_center = self.entity.rect.center()
            attr_center = attr.rect.center()
            
            # 计算连线的起点和终点
            entity_point = self._find_intersection_point(
                attr_center.x(), attr_center.y(),
                entity_center.x(), entity_center.y(),
                self.entity.rect, False
            )
            
            attr_point = self._find_intersection_point(
                entity_center.x(), entity_center.y(),
                attr_center.x(), attr_center.y(),
                attr.rect, True
            )
            
            # 绘制连线
            if attr.is_foreign:
                painter.setPen(QPen(Qt.black, 1, Qt.DashLine))
            else:
                painter.setPen(QPen(Qt.black, 1))
            
            painter.drawLine(
                int(entity_point[0]), int(entity_point[1]),
                int(attr_point[0]), int(attr_point[1])
            )
            
            # 绘制属性
            painter.setPen(QPen(Qt.black, 2))
            painter.drawEllipse(attr.rect)
            painter.setFont(attr_font)
            painter.drawText(attr.rect, Qt.AlignCenter, attr.name)

    def export(self, file_path: str):
        # 计算所有图形的边界
        min_x = float('inf')
        min_y = float('inf')
        max_x = float('-inf')
        max_y = float('-inf')
        
        # 检查实体的边界
        if self.entity:
            min_x = min(min_x, self.entity.rect.left())
            min_y = min(min_y, self.entity.rect.top())
            max_x = max(max_x, self.entity.rect.right())
            max_y = max(max_y, self.entity.rect.bottom())
        
        # 检查所有属性的边界
        for attr in self.attributes:
            min_x = min(min_x, attr.rect.left())
            min_y = min(min_y, attr.rect.top())
            max_x = max(max_x, attr.rect.right())
            max_y = max(max_y, attr.rect.bottom())
        
        # 添加一些边距
        padding = 20
        min_x -= padding
        min_y -= padding
        max_x += padding
        max_y += padding
        
        # 计算尺寸
        width = int(max_x - min_x)
        height = int(max_y - min_y)
        
        # 创建QPixmap
        pixmap = QPixmap(width, height)
        pixmap.fill(Qt.white)
        
        # 创建画笔
        painter = QPainter(pixmap)
        
        # 保存当前状态
        old_offset_x = self.offset_x
        old_offset_y = self.offset_y
        old_positions = []
        
        # 保存所有图形的原始位置
        for shape in [self.entity] + self.attributes:
            if shape:
                old_positions.append((shape, shape.rect.center()))
        
        # 移动所有图形到新坐标系
        for shape, _ in old_positions:
            if shape:
                old_center = shape.rect.center()
                new_x = old_center.x() - min_x
                new_y = old_center.y() - min_y
                shape.rect.moveCenter(QPointF(new_x, new_y))
        
        # 使用专门的导出绘制方法
        self.paint_for_export(painter, min_x, min_y, width, height)
        
        # 结束绘制
        painter.end()
        
        # 恢复所有图形的原始位置
        for shape, old_center in old_positions:
            if shape:
                shape.rect.moveCenter(old_center)
        
        # 恢复原始状态
        self.offset_x = old_offset_x
        self.offset_y = old_offset_y
        
        # 保存图像
        pixmap.save(file_path)
    
    def zoom_in(self):
        if self.scale < self.max_scale:
            self.scale *= 1.2
            self.widget.update()
    
    def zoom_out(self):
        if self.scale > self.min_scale:
            self.scale *= 0.8
            self.widget.update()
    
    def start_pan(self, pos):
        self.is_panning = True
        self.last_mouse_pos = pos
    
    def pan_to(self, pos):
        if self.is_panning and self.last_mouse_pos:
            dx = pos.x() - self.last_mouse_pos.x()
            dy = pos.y() - self.last_mouse_pos.y()
            self.offset_x = dx
            self.offset_y = dy
            self.last_mouse_pos = pos
            
            # 直接移动所有图形，而不是重新计算位置
            for shape in [self.entity] + self.attributes:
                if shape:
                    old_center = shape.rect.center()
                    new_x = old_center.x() + dx
                    new_y = old_center.y() + dy
                    shape.rect.moveCenter(QPointF(new_x, new_y))
            
            self.widget.update()
    
    def stop_pan(self):
        self.is_panning = False
        self.last_mouse_pos = None

    def paint(self, painter: QPainter):
        if not self.entity:
            return
        
        painter.setRenderHint(QPainter.Antialiasing)
        
        # 应用缩放变换
        painter.save()
        painter.translate(self.widget.width()/2, self.widget.height()/2)
        painter.scale(self.scale, self.scale)
        painter.translate(-self.widget.width()/2, -self.widget.height()/2)
        
        # 设置字体
        entity_font = QFont("Arial", 12, QFont.Bold)
        attr_font = QFont("Arial", 10)
        
        # 绘制实体
        painter.setPen(QPen(Qt.black, 2))  # 实体边框加粗
        painter.drawRect(self.entity.rect)
        painter.setFont(entity_font)
        painter.drawText(self.entity.rect, Qt.AlignCenter, self.entity.name)
        
        # 绘制连线和属性
        for attr in self.attributes:
            # 获取实体和属性的中心点
            entity_center = self.entity.rect.center()
            attr_center = attr.rect.center()
            
            # 计算连线的起点和终点
            entity_point = self._find_intersection_point(
                attr_center.x(), attr_center.y(),
                entity_center.x(), entity_center.y(),
                self.entity.rect, False
            )
            
            attr_point = self._find_intersection_point(
                entity_center.x(), entity_center.y(),
                attr_center.x(), attr_center.y(),
                attr.rect, True
            )
            
            # 绘制连线
            if attr.is_foreign:
                painter.setPen(QPen(Qt.black, 1, Qt.DashLine))
            else:
                painter.setPen(QPen(Qt.black, 1))
            
            painter.drawLine(
                int(entity_point[0]), int(entity_point[1]),
                int(attr_point[0]), int(attr_point[1])
            )
            
            # 绘制属性，所有属性边框都加粗
            painter.setPen(QPen(Qt.black, 2))  # 属性边框加粗
            painter.drawEllipse(attr.rect)
            painter.setFont(attr_font)
            painter.drawText(attr.rect, Qt.AlignCenter, attr.name)
        
        # 恢复变换
        painter.restore()